/**
* The deformation code. This file deals with setting up and solving
* the matrix system.
*/


#include <OGF/quick_start/common/deform.h>
#include <OGF/quick_start/common/select_util.h>
#include <OGF/basic/os/stopwatch.h>
#include <OGF/cells/map/map_cell_heap.h>

#include <set>
#include <stack>
#include <queue>
#include <limits>

#include <OGF/math/linear_algebra/vector.h>

namespace OGF 
{
	/**
	*  A 3 rows to the neighbour transformation matrix which represent
	*  the transformation on this vertex
	*  [  Vkx     0   Vkz  -Vky  1  0  0]
	*  [  Vky  -Vkz     0   Vkx  0  1  0]
	*  [  Vkz   Vky  -Vkx     0  0  0  1]
	*/
	void Deform::fill_neighbour_transform_rows(Map::Vertex *vert, LinearSolver *T_solver){
		double kx, ky, kz, new_kx, new_ky, new_kz;

		// Transform relative to initial/baked position.
		kx = (*true_pos)[vert].x();
		ky = (*true_pos)[vert].y();
		kz = (*true_pos)[vert].z();

		// Where the transformation has taken us.
		new_kx = vert->point().x();
		new_ky = vert->point().y();
		new_kz = vert->point().z();
		//new_kx = (verts[id])->point()[0];
		//new_ky = (verts[id])->point()[1];
		//new_kz = (verts[id])->point()[2];

		// Form three rows for this vertex's neighbours
		T_solver->begin_row();
		T_solver->add_coefficient(0,kx);
		T_solver->add_coefficient(2,kz);
		T_solver->add_coefficient(3,-1*ky);
		T_solver->add_coefficient(4,1.0f);
		T_solver->set_right_hand_side(new_kx);
		T_solver->end_row();

		T_solver->begin_row();
		T_solver->add_coefficient(0,ky);
		T_solver->add_coefficient(1,-1*kz);
		T_solver->add_coefficient(3,kx);
		T_solver->add_coefficient(5,1.0f);
		T_solver->set_right_hand_side(new_ky);
		T_solver->end_row();

		T_solver->begin_row();
		T_solver->add_coefficient(0,kz);
		T_solver->add_coefficient(1,ky);
		T_solver->add_coefficient(2,-1*kx);
		T_solver->add_coefficient(6,1.0f);
		T_solver->set_right_hand_side(new_kz);
		T_solver->end_row();
	}

	/**
	*  Transform all the delta-vectors in the ROI, according to the best fit
	*  linear transformation of the local neighbourhood. Allow translation,
	*  rotation and uniform scale transformations.
	*  Precondition: verts stores the current vertex positions
	*		 true_verts stores the undeformed vertex positions
	*		 delta stores the current delta vectors
	*		 true delta stores the undeformed delta vectors
	*  Postcondition: deltas are transformed optimally based on transformation 
	*		  of local neighbourhood.
	*/
	void Deform::optimally_transform_deltas()
	{
		std::cout<<" Optimally transforming Laplacian delta vectors"<<std::endl;

		SystemSolverParameters params ;
       	params.set_arg("method", "SUPERLU") ;
		
		//FOR_EACH_VERTEX(Map, surface, it)
		//{
		//	Map::Vertex *v = it;
		for(unsigned int i=0;i<verts.size();i++)
		{
			Map::Vertex *v = verts[i];

			// Only do the transformation of a vertex if it is going to have a laplace row in the solver.
			int clas = (*classify)[v];
			bool has_laplace_row = (clas == GUARD || clas == CORE || clas == ANCHOR);
			if(has_laplace_row)
			{
				// Find the optimal transformation Ti for each vertex in the region of interest
				// Allow only rotations, uniform scales, and translations
				LinearSolver T_solver(7);
				T_solver.set_system_solver(params) ;
				T_solver.set_least_squares(true) ;
				T_solver.set_invert_matrix(true) ;
				// symmetric is TRUE for least squares systems!!! A^TA 
				// and is set automatically when you set_least_squares
				//T_solver.set_symmetric(false); 
				T_solver.set_quiet(true);
				T_solver.begin_system() ;
				Map::Vertex *adjacent;

				//Lock the scale, so we don't grow/shrink
				//T_solver.variable(0).set_value(1);
				//T_solver.variable(0).lock();

				FOR_EACH_ADJACENT_VERTEX(v, adjacent)
				{
					// Adjacent should never be outside the ROI, 
					// but even if it is, it doesn't matter. It just
					// means that we get a slightly different transform.
					fill_neighbour_transform_rows(adjacent,&T_solver);
				}
				fill_neighbour_transform_rows(v,&T_solver);

				T_solver.end_system();
				T_solver.solve() ;

				// Form the transformation matrix from the solution
				// Ignore translations
				//double s  = T_solver.variable(0).value();
				double s = 1.;
				double h1 = T_solver.variable(1).value();
				double h2 = T_solver.variable(2).value(); 
				double h3 = T_solver.variable(3).value();

				Matrix2d m;
				m(0,0) =     s ; m(0,1) =   -h3 ; m(0,2) =    h2 ; 
				m(1,0) =    h3 ; m(1,1) =     s ; m(1,2) =   -h1 ;
				m(2,0) =   -h2 ; m(2,1) =    h1 ; m(2,2) =     s ; 

				// Apply transformation: update the deltas
				Vector3d delta = (*true_deltas)[ v ];
				delta = mult(m, delta);
				(*deltas)[v] = delta;

			}
		}
	}

	/**
	* Setup the solver object if it is not ready.
	* pre: "verts" has been built, as by populate_verts()
	*/
	void Deform::build_system_solver()
	{
		if(solver != nil)
			return;

        SystemSolverParameters params ;
        params.set_arg("method", "SUPERLU") ; // "CG" is WAY slower

		std::cout<<" Constructing solver object"<<std::endl;
		solver = new LinearSolver(verts.size()) ;
		solver->set_system_solver(params) ;
		solver->set_least_squares(true) ;
		solver->set_invert_matrix(true) ;
	}

	/**
	* Just do the linear laplacian part, for fast interaction with the GUI.
	*/
	void Deform::do_fast_deformation()
	{
		if( !is_anchor_good() )
		{
			std::cout<< "Bad anchor, skipping deformation" <<std::endl;
			return;
		}
		(*constraint_pos)[anchor] = anchor->point();

		std::cout<<"\n\n==== Starting deformation ===="<<std::endl;
        SystemStopwatch w ;

		populate_verts();
		std::cout<<"Elapsed time " << w.elapsed_user_time()<<" s"<<std::endl ;
		
		build_system_solver();

		do_laplace_solve();
		std::cout<<"Elapsed time " << w.elapsed_user_time()<<" s"<<std::endl ;

		std::cout<<"==== Deformation complete ====\n\n"<<std::endl;
	}

	/**
	* Do the deformation.
	* pre:
	*   - ROI is calculated by grow_region()
	*/
	void Deform::do_deformation()
	{
		if( !is_anchor_good() )
		{
			std::cout<< "Bad anchor, skipping deformation" <<std::endl;
			return;
		}
		(*constraint_pos)[anchor] = anchor->point();

		std::cout<<"\n\n==== Starting deformation ===="<<std::endl;
        SystemStopwatch w ;

		populate_verts();
		std::cout<<"Elapsed time " << w.elapsed_user_time()<<" s"<<std::endl ;
		
		build_system_solver();

		//Reset the deltas to the originals, and do a linear solve.
		FOR_EACH_VERTEX(Map, surface, v)
			(*deltas)[ v ] = (*true_deltas)[ v ];
		do_laplace_solve();
		std::cout<<"Elapsed time " << w.elapsed_user_time()<<" s"<<std::endl ;

		for(int i=0;i<1;i++)
		{
			// Update the deltas to match the transformation to their neighbourhood
			optimally_transform_deltas();
			std::cout<<"Elapsed time " << w.elapsed_user_time()<<" s"<<std::endl ;

			// Re-solve, with the new deltas
			do_laplace_solve();
			std::cout<<"Elapsed time " << w.elapsed_user_time()<<" s"<<std::endl ;
		}

		std::cout<<"==== Deformation complete ====\n\n"<<std::endl;

	}

	void Deform::do_laplace_solve()
	{
		static char coord_name[4] = "xyz";

		//for each coordinate:
		for(int coord=0; coord<3; coord++)
		{
			std::cout<<" Solving for "<<coord_name[coord]<<" coordinates."<<std::endl;

			//std::cout<<"  Putting current positions into solver"<<std::endl;
			verts_to_solver(coord);

            solver->begin_system() ;
            setup_Laplace(coord) ;
            solver->end_system() ;
			
			//std::cout<<"Elapsed time " << w.elapsed_user_time()<<"s"<<std::endl ;
			std::cout<<"  Solving least-squares system"<<std::endl;
            solver->solve() ;

			std::cout<<"  Applying results to mesh"<<std::endl;
            solver_to_verts(coord) ;

			solver->restart();
			//std::cout<<"Elapsed time " << w.elapsed_user_time()<<"s"<<std::endl ;
		}
	}

	bool Deform::has_laplace_row(Map::Vertex *v)
	{
		int vert_class = (*classify)[v];
		return vert_class == GUARD || vert_class == ANCHOR || vert_class == CORE;
	}
	bool Deform::has_constraint_row(Map::Vertex *v)
	{
		int vert_class = (*classify)[v];
		return vert_class == GUARD;// || vert_class == ANCHOR || vert_class == BORDER;
	}


	/**
	* Fill the solver with the laplace matrix, and 
	* component 'coord' of the delta's/positions
	*/
	void Deform::setup_Laplace(
		int coord)
	{
		std::cout<<"  Constructing Laplacian matrix problem. Rows & RHS"<<std::endl;

		// The Laplace Matrix.
		unsigned int n_verts = verts.size();
		unsigned int row = 0;
		for(unsigned int i=0;i<n_verts;i++)
		{
			Map::Vertex *v = verts[i];
			if( has_laplace_row(v) )
			{
				// Each vertex MUST be on the right row. It's '1' entry
				// must be on the diagonal.
				// This is only requried because the transpose assumes it.
				ogf_assert( row == (*index)[ v ]);
				row++;

				fill_laplace_row(v, coord);
			}
		}
		for(unsigned int i=0;i<n_verts;i++)
		{
			if( has_constraint_row(verts[i]) )
				fill_constraint_row(verts[i], coord);
		}
		/*
		for(unsigned int i=0;i<n_verts;i++)
		{
			Map::Vertex *v = verts[i];
			if ( has_laplace_row(v) )
			{
				// Fill the spring rows for all of my neighbours
				/*MapVertexAttribute<double> stiffness(surface, "stiffness");
				double weight = 0.1 + 2.*stiffness[v];
				Map::Vertex *adjacent;
				
				FOR_EACH_ADJACENT_VERTEX(v, adjacent)
				{
					// insert a row for this neighbour
					solver->begin_row();
					solver->add_coefficient((*index)[v], weight);
					solver->add_coefficient((*index)[adjacent],-weight);
					solver->set_right_hand_side((weight*((*true_pos)[v][coord] - (*true_pos)[adjacent][coord])));
					solver->end_row();
				}*/
				/*
				// insert a pseudo-gravity row
				if(coord == GRAVITY_COORD){
					double mymass = (*mass)[v];
					solver->begin_row();
					solver->add_coefficient((*index)[v], mymass);
					solver->set_right_hand_side(mymass*(*true_pos)[v][coord]);
					solver->end_row();
				}
			}
		}*/

		// This code works, and is fast, but is harder to make agree with the transpose:
		//fill_laplace_row(anchor, coord);
		//for(unsigned int i=0;i<verts_core.size();i++)
		//	fill_laplace_row(verts_core[i], coord);
		//for(unsigned int i=0;i<verts_guard.size();i++)
		//	fill_laplace_row(verts_guard[i], coord);
		// The constraints.
		//for(unsigned int i=0;i<verts_guard.size();i++)
		//	fill_constraint_row(verts_guard[i], coord);
		//for(unsigned int i=0;i<verts_border.size();i++)
		//	fill_constraint_row(verts_border[i], coord);
		//fill_constraint_row(anchor, coord);
	}

	/**
	* Add a row to the matrix to lock a vertex in place.
	* [0 .. 0 1 0 .. 0] x = x0
	*/
	void Deform::fill_constraint_row(
		Map::Vertex *vert, 
		int coord)
	{
		int row = (*index)[vert];
		//Constrait row is just a 1 on the diagonal,
		//and the RHS is just the position that I want
		
		solver->begin_row();
		solver->add_coefficient(row, constraint_weight);
		
		double u = constraint_weight * (*constraint_pos)[vert][coord];

		if(coord == GRAVITY_COORD)//gravity acts in this direction..
			u = u - gravity*alpha[row];

		solver->set_right_hand_side(u);
        solver->end_row() ;
	}

	void Deform::fill_laplace_row( Map::Vertex *vert, int coord )
	{
		// Regular row in the laplacian matrix has:
		//   +1 in column j, when this is vertex j
		//   -1/degree in each column j for which vertex j is a neighbor
		// Thus, the entire row adds up to zero.

		//double weight = (*laplace_weight)[vert];
		MapVertexAttribute<double> stiffness(surface, "stiffness");
		double weight = 1.;
		/*weight = weight * (1.0 + 1.0 * stiffness[vert]);
		if(weight != 1.0){
			std::cout << "weight is " << weight << std::endl;
		}*/
		float inv_degree = -1.0f/vert->degree();

		solver->begin_row();
		solver->add_coefficient((*index)[vert], weight);

		Map::Vertex *adjacent;
		FOR_EACH_ADJACENT_VERTEX(vert, adjacent)
		{
			int column = (*index)[adjacent];
			solver->add_coefficient(column, inv_degree*weight);
		}

		Vector3d delta_vec = (*deltas)[vert]; //remember, this structure holds the rotated deltas
		double rhs = delta_vec[coord]*weight;

		if(coord == GRAVITY_COORD)//gravity acts in this direction..
			rhs = rhs - gravity*alpha[(*index)[vert]];

		solver->set_right_hand_side(rhs);

        solver->end_row() ;
	}

	/**
	* Takes the current shape of the mesh, and uses it as the
	* base for future deformations. 
	* Specifically, this calculates and stores:
	*    - delta vectors
	*    - positions of verticies
	*/
	void Deform::bake_deformation()	//Saves the current mesh shape
	{
		this->compute_deltas();
		FOR_EACH_VERTEX(Map, surface, v)
			(*true_pos)[v] = v->point();
	}

	/**
	* Compute the Laplacian deltas of all points in the mesh,
	* and store them in the 'true_deltas' attribute.
	* We only want to do this on the original (un-deformed) mesh,
	* as that detail information is what we want to preserve.
	*/
	void Deform::compute_deltas()
	{
		std::cout<<"Computing delta vectors"<<std::endl;
		FOR_EACH_VERTEX(Map, surface, it)
		{
			Map::Vertex *vert = &(*it);

			double px, py, pz; //predicted position, average of neighbors
			px = py = pz = 0;
			
			Map::Vertex *v;
			int degree = 0;
			FOR_EACH_ADJACENT_VERTEX(vert, v)
			{
				Point3d pos = v->point();

				px += pos[0];
				py += pos[1];
				pz += pos[2];

				degree++;
			}

			Point3d avg(px/degree, py/degree, pz/degree);
			(*deltas)[vert] = (*true_deltas)[vert] = vert->point() - avg;

		}
	}

	/**
	* Put all the vertex information into the solver.
	* Putting in the coordinates is only really needed for iterative solvers,
	* the important point is locking all the right variables.
	*/
	void Deform::verts_to_solver(int coord)
	{
		//std::cout<<"Putting initial "<<verts.size()<<" vertex coords into solver"<<std::endl;
		for(unsigned int i=0;i<verts.size();i++)
		{
			Map::Vertex *v = verts[i];

			//Change of heart, not locking anything
			int v_class = (*classify)[v];
			if(v_class == ANCHOR || v_class == BORDER)
				solver->variable(i).lock();

			double r = v->point()[coord];
			solver->variable(i).set_value(r);
		}
	}
	/**
	* Read the solution from the solver into the mesh.
	*/
	void Deform::solver_to_verts(int coord)
	{
		for(unsigned int i=0;i<verts.size();i++)
		{
			double r = solver->variable(i).value();
			verts[i]->point()[coord] = r;
		}
	}	
	/**
	* For every vertex in the ROI, get the length (cartesian) of the shortest edge-path
	* to a vertex in the argument vector. Basically, a modified dijkstra's algorithm.
	*/
	/*void Deform::compute_geodesic_field(std::vector<Map::Vertex*> seed, MapVertexAttribute<double> *distance)
	{
		ogf_assert(distance != nil);

		MapVertexHeap queue(surface, *distance);

		//MapVertexAttribute<bool> is_visited(surface);
		FOR_EACH_VERTEX(Map, surface, it)
		{
			(*distance)[ it ] = 1e19; //infinity
		}

		//std::cout<<"Initializing distance on seed verticies"<<std::endl;
		for(unsigned int i=0;i<seed.size();i++)
		{
			Map::Vertex *vert = seed[i];
			ogf_assert(vert != nil);
			(*distance)[ vert ] = 0;
			queue.push( vert );
		}
		
		//std::cout<<"Done init, now recursing."<<std::endl;
		while(!queue.empty()) //while there are more verticies in the spreading front.
		{
			Map::Vertex* vert = queue.pop();
			//std::cout<<"Visiting vertex"<<std::endl;	

			double distance_so_far = (*distance)[ vert ];

			//Add all this vert's adjacent verticies to the queue to inspect.
			Map::Halfedge* first_edge = vert->halfedge();
			Map::Halfedge* edge = first_edge;
			do
			{
				Map::Vertex* adj_vert = edge->opposite()->vertex();

				bool in_region = is_in_region(adj_vert);

				if(in_region)
				{
					Vector3d edge_vec = vert->point() - adj_vert->point();
					double edge_length = edge_vec.norm();

					double new_dist = distance_so_far + edge_length;
					
					double old_dist = (*distance)[ adj_vert ];
					//double first_visit = !is_visited[ adj_vert ];
					if( new_dist <  old_dist ) //If we've never visted before, OR this route was shorted than the last one
					{
						(*distance)[ adj_vert ] = new_dist;
						//is_visited[ vert ] = true;
						//std::cout<<"dist = "<<new_dist<<std::endl;
						if(queue.contains(adj_vert))
							queue.update_cost(adj_vert);
						else
							queue.push(adj_vert);
					}

				}
				
				edge = edge->next_around_vertex();
			}while(edge != first_edge);
		}
	}*/

	/**
	* Compute the geodesic distances to anchor and boundary.
	*/
/*	void Deform::compute_geodesics()
	{
		std::cout<<" Computing geodesic distances to anchor and boundary"<<std::endl;
		std::vector<Map::Vertex*> anchor_vec;
		anchor_vec.push_back(anchor);
		this->compute_geodesic_field(anchor_vec, distance_to_anchor);
		this->compute_geodesic_field(verts_border, distance_to_border);
	}*/
	/**
	* Rotate all the delta-vectors in the ROI, according to an interpolation
	* between the anchor angle and zero at the boundary. Interplate using the
	* geodesic distances.
	* precondition: geodesic_distances() has been called for this ROI
	*/
/*	void Deform::rotate_deltas()
	{
		//rotation angle is now:
		// angle = angle*distance_to_border/(dist_to_anchor+distance_to_border);

		std::cout<<" Propogating anchor rotation ("<<this->angle<<" radians)"<<std::endl;
		std::cout<<" Rotating Laplacian delta vectors"<<std::endl;
		MapTexVertexNormal normal(surface) ;
		FOR_EACH_VERTEX(Map, surface, it)
		{
			Map::Vertex *v = it;
			(*rotation_angle)[ v ] = 0;
			if(this->is_in_region(v))
			{
				double a = (*distance_to_anchor)[v];
				double b = (*distance_to_border)[v];
				double weight = b/(a+b);
				double angle  = weight * this->angle;
				(*rotation_angle)[ v ] = angle;

				Matrix2d m = this->fromAxisAngle(this->axis, angle);
				Vector3d delta = (*true_deltas)[ v ];
				(*deltas)[v] = mult(m, delta );
			}
		}
	}*/

} //namespace
